import time
import torch
from cat_dog_torch.settings import *
import torch.nn.functional as F
from torch.optim import Adam
from cat_dog_torch.model import AlexNet
from cat_dog_torch.loaders import train_dataloader, val_dataloader

DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
start_time = time.time()

def compute_acc(model, dataloader, device):
    correct_pred, num_examples = 0, 0  # correct_pred 统计正确预测的样本数，num_examples 统计样本总数
    for i, (features, targets) in enumerate(dataloader):
        features = features.to(device)
        targets = targets.to(device)

        logits = model(features)
        _, predicted_labels = torch.max(logits, 1)

        num_examples += targets.size(0)
        assert predicted_labels.size() == targets.size()
        correct_pred += (predicted_labels == targets).sum()

    return correct_pred.float() / num_examples * 100


if __name__ == '__main__':

    model = AlexNet()
    model = model.to(DEVICE)

    cost_list = []
    train_acc_list, val_acc_list = [], []
    start_time = time.time()

    optimizer = Adam(model.parameters(), lr = RATE)

    for epoch in range(EPOCHS):
        model.train()
        for batch_idx, (features, targets) in enumerate(train_dataloader):
            features = features.to(DEVICE)
            targets = targets.to(DEVICE)

            optimizer.zero_grad()
            logits = model(features)
            cost = F.cross_entropy(logits, targets)
            cost.backward()

            optimizer.step()
            cost_list.append(cost.item())

            if not batch_idx % 50:
                print('epoch: %03d/%03d | Batch %03d/%03d | cost: %.4f' %
                      (epoch + 1,  EPOCHS, batch_idx, len(train_dataloader), cost.float()))

        model.eval()
        with torch.set_grad_enabled(False):
            train_acc = compute_acc(model, train_dataloader, device=DEVICE)
            val_acc = compute_acc(model, val_dataloader, device=DEVICE)

            print('Epoch: %03d/%03d | Training ACC: %.4f%% | Validation ACC: %.4f%%' %
                  (epoch + 1, EPOCHS, train_acc, val_acc))

    print('Total Time Elapsed: %.2f min' % ((time.time() - start_time) / 60))
    torch.save(model.state_dict(), MODEL_NAME)